#pragma once

#include "fsa_shared_field.hpp"
#include "accept_char.hpp"
#include <iostream>
#include <algorithm>
#include <cmath>
#include <map>
#include <memory>
#include <set>
#include <string>
#include <unordered_map>
#include <cmath>
#include <optional>

using std::map;
using std::set;
using std::string;
using std::unordered_map;

#define __FSA_VERIFY_SINGLE_AC(fsa)                                            \
    if ((fsa)._acStates.size() != 1){\
        std::cerr << "FSA has multiple accept states\n";                            \
        throw std::range_error("FSA has multiple accept states");              \
    }

#define __CHECK_SAME_SHARED(fsa)                                               \
    if (S != fsa.S) {                                                          \
        std::cerr << "FSA is not based on same shared field!\n";                    \
        throw std::runtime_error("FSA is not based on same shared field!");    \
    }

/**
 * Finite State Automata
 */
class FSA {
private:
    int _beginState;
    const std::pair<const string, const int> _TOKENINFO;
    const int _STATE_BASE_POS; // _states[stateBasePos, -1] belongs to this FSA
    // <AC state index, <token-name, priority>>, smaller value means higher priority
    map<int, std::pair<const string, const unsigned int>> _acStates;
    FsaSharedField *S;
    enum class FsaType { EPSILON, NFA, DFA, MIN };
    FsaType fsaType;
    FSA(FsaSharedField *ss)
        : _TOKENINFO({"This is the base FSA.", 999999}), _STATE_BASE_POS(0) {
        S = ss;
        // S->newState(); // reserved for base FSA because of a stupid bug
        _beginState = 0;
        fsaType = FsaType::MIN;
    }

public:
    static FSA makeBase(FsaSharedField *ss) {
        if (ss->stateSize() != 0) {
            throw std::runtime_error("Shared field already has base FSA!");
        }
        return FSA(ss);
    }
    FSA(FsaSharedField *ss, const string &tokenName, const int priority)
        : _STATE_BASE_POS(ss->stateSize()), _TOKENINFO({tokenName, priority}) {
        S = ss;
        fsaType = FsaType::NFA;
        _beginState = S->stateSize();
        _acStates.insert({_beginState, _TOKENINFO});
        S->newState();
    }

    /****************************************************************
     * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * *
     * * * * * * * * * * * * Regex Operations * * * * * * * * * * * *
     * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * *
     ****************************************************************
     */

    // wrapper for starOp(), plusOp(), questionOp90
    void performOp(char op, bool frombegin) {
        if (op == '*') {
            starOp(frombegin);
        } else if (op == '+') {
            plusOp(frombegin);
        } else if (op == '?') {
            questionOp(frombegin);
        } else {
            std::cerr << "Illegal operation!";
            throw std::runtime_error("Illegal operation!");
        }
    }
    void starOp(bool frombegin) {
        __FSA_VERIFY_SINGLE_AC(*this);
        fsaType = FsaType::EPSILON;
        plusOp(frombegin);
        questionOp(frombegin);
    }
    void orOp(const FSA &rhs) {
        __CHECK_SAME_SHARED(rhs);
        __FSA_VERIFY_SINGLE_AC(*this);
        __FSA_VERIFY_SINGLE_AC(rhs);
        fsaType = FsaType::EPSILON;
        const int oldBegin = _beginState;
        _beginState = S->stateSize();
        // ====================Process start==================================
        S->newState();
        S->cpAndMergeState(rhs._beginState, _beginState);
        S->cpAndMergeState(oldBegin, _beginState);
        // ====================Process AC=====================================
        const int oldAc = _acStates.begin()->first;
        _acStates.clear();
        _acStates.insert({S->stateSize(), _TOKENINFO}); // new ac state
        S->newState();
        S->addTransition(oldAc, _acStates.begin()->first, {__ID_EPSILON});
        S->addTransition(rhs._acStates.begin()->first, _acStates.begin()->first,
                         {__ID_EPSILON});
    }
    void concatOp(const FSA &rhs) {
        __CHECK_SAME_SHARED(rhs);
        __FSA_VERIFY_SINGLE_AC(*this);
        fsaType = std::min(fsaType, rhs.fsaType);
        const int endState = maxKey(_acStates);
        if (endState - _beginState == 0) { // lhs is empty
            S->addTransition(_acStates.begin()->first, rhs._beginState,
                             {__ID_EPSILON});
        } else {
            for (int state = _beginState; state <= endState; ++state) {
                S->redirect(state, _acStates.begin()->first, rhs._beginState);
            }
        }
        S->cpAndMergeState(_acStates.begin()->first, rhs._beginState);
        _acStates = rhs._acStates;
    }
    void plusOp(bool frombegin) {
        __FSA_VERIFY_SINGLE_AC(*this);
        if(frombegin){
            S->cpAndMergeState(_beginState, S->stateSize() - 1);
        } else {
            for (int sym = 0; sym < S->symSize(); ++sym){
                if (S->stateAt(S->stateSize() - 2)[sym].count(S->stateSize() - 1)) {
                    S->stateAt(S->stateSize() - 1)[sym].insert(S->stateSize() - 1);
                }
            }
        }
    }
    void questionOp(bool frombegin) {
        __FSA_VERIFY_SINGLE_AC(*this);
        fsaType = FsaType::EPSILON;
        if(frombegin){
            S->addTransition(_beginState, _acStates.begin()->first, {__ID_EPSILON});
        }else{
            S->addTransition(S->stateSize() - 2, _acStates.begin()->first, {__ID_EPSILON});
        }
    }

    /**
     * append a state to the end of current FSA
     */
    void appendState(acceptableChars s) {
        __FSA_VERIFY_SINGLE_AC(*this);
        const int fromStateId = S->stateSize() - 1;
        const int toStateId = S->stateSize();
        _acStates.clear();
        _acStates.insert({toStateId, _TOKENINFO}); // new ac state
        S->newState();
        set<int> symbols; // symbol idx to accept this transition
        for (int sym = 0; sym < S->symSize(); ++sym) { // compress symbol set
            if (s.includes(S->symAt(sym))) {
                symbols.insert(sym);
                s.setDifference(S->symAt(sym));
                if (s == S->symAt(sym))
                    break;
            } else {
                acceptableChars intersection = s.getIntersection(S->symAt(sym));
                if (!intersection.empty()) {
                    // sym[i] = sym[i] - s
                    S->symAt(sym).setDifference(s);
                    // intersection is the new symbol
                    const int newSymbolId = S->symSize();
                    symbols.insert(newSymbolId);
                    S->newSymbol(intersection);
                    // update states affected by splitting sym[i]
                    for (int state = 0; state < S->stateSize(); ++state) {
                        if (!S->stateAt(state)[sym].empty()) {
                            S->stateAt(state)[newSymbolId] = S->stateAt(state)[sym];
                        }
                    }
                    // s = s - sym[i]
                    s.setDifference(S->symAt(sym));
                } else { // ignored;
                }
            }
        }
        if (!s.empty()) {
            symbols.insert(S->symSize());
            S->newSymbol(s);
        }
        S->addTransition(fromStateId, toStateId, symbols);
    }

    /****************************************************************
     * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * *
     * * * * * * * * * * * * Convert to DFA * * * * * * * * * * * * *
     * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * *
     ****************************************************************
     */
public:
    void convertDFA(const std::pair<const string, const int> &tokenInfo) {
        fsaType = FsaType::EPSILON;
        printFSA("Epsilon");
        if (fsaType == FsaType::EPSILON) {
            removeEpsilonEdge(tokenInfo);
        }
        printFSA("NFA");
        if (fsaType == FsaType::EPSILON || fsaType == FsaType::NFA) {
            toDfaInternal();
            removeNfaStates();
        }
        printFSA("DFA");
        if (fsaType == FsaType::DFA) {
            minDfa();
        }
    }

private:

    /**
     * Computes epsilon closure of a state
     * @return <symbol, <states>> map
     */
    unordered_map<int, set<int>>
    epsilonClosure(int from, int state, bool **vis,
                   const std::pair<const string, const int> &tokenInfo) {
        if(vis[from][state])
            return {};
        vis[from][state] = true;
        unordered_map<int, set<int>> ret;
        for(auto i: S->stateAt(state)[0]){
            auto tmp = epsilonClosure(state, i, vis, tokenInfo);
            for(auto &[sym, tos]: tmp){
                ret[sym].insert(tos.begin(), tos.end());
            }
            if (_acStates.count(i)) {
                _acStates.insert({state, tokenInfo});
            }
        }
        for (int sym = 1; sym < S->symSize(); ++sym) {
            if(!S->stateAt(state)[sym].empty()){
                ret[sym].insert(S->stateAt(state)[sym].begin(),
                                S->stateAt(state)[sym].end());
            }
        }
        return ret;
    }

    void removeEpsilonEdge(const std::pair<const string, const int> &tokenInfo) {
        fsaType = FsaType::NFA;
        bool **vis = new bool*[S->stateSize()];
        //TODO is S->stateSize() correct??? or its keyMax(acStates)
        for (int state = _beginState; state < S->stateSize(); ++state){
            S->stateAt(state)[0].erase(state);
            //FIXME due to a stupid behaviour in concat
        }
        for (int state = _beginState; state < S->stateSize(); ++state){
            // initialize visit
            for (int i = 0; i < S->stateSize(); ++i) {
                vis[i] = new bool[S->stateSize()];
                for (int j = 0; j < S->stateSize(); ++j){
                    vis[i][j] = false;
                }
            }
            auto epsilons = S->stateAt(state)[0];
            auto closure = epsilonClosure(state, state, vis, tokenInfo);
            S->stateAt(state).clear();
            S->stateAt(state).resize(S->symSize());
            for (auto &[sym, tos] : closure) { // reconstruct to sets
                S->stateAt(state)[sym] = tos;
            }
        }
    }

    void toDfaInternal() {
        fsaType = FsaType::DFA;
        map<set<int>, int> lookup; // <tos, stateId>
        const int oldBegin = _beginState;
        // NFA states at [oldBegin, beginState),
        // DFA states at [beginState, size)
        _beginState = S->stateSize(); // new begin!!!
        if(_acStates.count(oldBegin)){
            _acStates.insert({_beginState, _TOKENINFO});
        }
        S->newState();
        S->cpAndMergeState(oldBegin, _beginState);
        for (int i = _beginState; i < S->stateSize(); ++i) {
            for (auto &tos : S->stateAt(i)) { // for each symbol
                if (tos.empty())
                    continue;
                if (lookup.count(tos)) {
                    const int newStateId = lookup.at(tos);
                    tos.clear();
                    tos.insert(newStateId);
                } else {
                    const int newStateId = S->stateSize();
                    lookup[tos] = newStateId;
                    S->newState();
                    for (auto to : tos) { // for each target state
                        S->cpAndMergeState(to, newStateId);
                        if (_acStates.count(to)) { // update ac states
                            _acStates.insert({newStateId, _acStates.at(to)});
                            // insertion fails
                            if(_acStates.at(newStateId) != _acStates.at(to)){
#ifdef __DEBUG
                                std::cerr << "[Warn]: AC state (" << newStateId << ", "
                                          << _acStates.at(to).first << ") already exists, ";
#endif
                                // force update due to higher priority
                                if(_acStates.at(to).second < _acStates.at(newStateId).second) {
                                    _acStates.erase(newStateId);
                                    _acStates.insert({newStateId, _acStates.at(to)});
#ifdef __DEBUG
                                    std::cerr << "force update to (" << to << ", "
                                              << _acStates.at(to).first
                                              << ") due to higher priority";
                                } else {
                                    std::cerr << "fail to insert state (" << to
                                              << ", " << _acStates.at(to).first << ").\n";
#endif
                                }
                            }
                        }
                    }
                    tos.clear();
                    tos.insert(newStateId); // redirect to new set
                }
            }
        }
    }

    void removeNfaStates() {
#ifdef __DEBUG
        __PRINT_LINE_HINT("Warning: would erase S->_states");
#endif
        const int offset = _beginState - _STATE_BASE_POS;
        // 1. redirect tos
        for (int i = _beginState; i < S->stateSize(); ++i) {
            for (auto &tos : S->stateAt(i)) {
                set<int> tmp;
                for (auto to : tos) {
                    tmp.insert(to - offset);
                }
                tos = tmp;
            }
        }
        // 2. erase NFA states
        S->eraseStates(_STATE_BASE_POS, _beginState);
        // 3. redirect begin and ac states
        map<int, std::pair<const string, const unsigned int>> tmp;
        for (auto [i, tokenName] : _acStates) {
            if (i - _beginState >= 0) {
                tmp.insert({i - offset, tokenName});
            }
        }
        _acStates = tmp;
        _beginState = _STATE_BASE_POS;
    }

    void minDfa() {
        fsaType = FsaType::MIN;
        // TODO minDFA
    }

public:
    void printFSA(const string &hint) {
#ifdef __DEBUG_LV1
        std::cout << std::endl << hint << std::endl;
        std::cout << "Begin state: " << _beginState << std::endl;
        std::cout << "AC states: ";
        for (auto [i, name] : _acStates) {
            std::cout << "(" << i << ", " << name.first << "); ";
        }
        std::cout << std::endl;
        S->printStates();
        std::cout << std::endl;
#endif
    }

    /****************************************************************
     * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * *
     * * * * * * * * Merge begin with another FSA * * * * * * * * * *
     * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * *
     ****************************************************************
     */
public:
    void mergeBegin(const FSA &rhs) {
        __CHECK_SAME_SHARED(rhs);
        if (_beginState != 0) {
            throw std::runtime_error("lhs is not base FSA");
        // } else if (maxKey(rhs._acStates) != S->stateSize() - 1) {
            // throw std::runtime_error("rhs is not at stack top");
        } else if (fsaType != FsaType::DFA && fsaType != FsaType::MIN) {
            throw std::runtime_error("lhs is not DFA");
        } else if (rhs.fsaType != FsaType::DFA && rhs.fsaType != FsaType::MIN) {
            throw std::runtime_error("rhs is not DFA");
        }
        _acStates.insert(rhs._acStates.begin(), rhs._acStates.end());
        S->cpAndMergeState(rhs._beginState, 0);
        S->redirect(rhs._STATE_BASE_POS, rhs._beginState, 0); // redirect to begin
        for (auto i : S->stateAt(_beginState)) {
            if (i.size() > 1) {
                fsaType = FsaType::NFA;
                break;
            }
        }
        convertDFA(rhs._TOKENINFO);
    }

    std::pair<const string, const int> getTokenInfo() { return _TOKENINFO; }
    int getBeginState() const { return _beginState; }
    std::optional<string> getAcName(int state) const {
        if(_acStates.count(state)){
            return _acStates.at(state).first;
        } else {
            return std::nullopt;
        }
    }
    const map<int, const string> getAcStates() const { 
        map<int, const string> ret;
        for(auto [key, val]: _acStates){
            ret.insert({key, val.first});
        }
        return ret;
    }
};
